import copy
from Crypto.Hash import SHA256
from numpy import mod
from train_model import TrainModel
import net

class Client:
    def __init__(self, client_idx, client_train_data, client_test_data, client_data_num, model, psi = False, pipe_conn = None, device = 'cuda') -> None:
        self.client_idx = client_idx
        self.train_data = client_train_data
        self.test_data = client_test_data
        self.train_num, self.test_num = client_data_num
        self.device = device
        self.model = copy.deepcopy(model)
        if psi == True:
            self.intersection = []
            self.each_label_num_dict = {}
            self.label_list = []
            self.hash_label_list = []
            self.pipe_conn = pipe_conn

    def update_local_dataset(self, client_idx, train_data, test_data, data_num,):
            self.client_idx = client_idx
            self.train_data = train_data
            self.test_data = test_data
            self.train_num, self.test_num = data_num

    def update_model(self, model):
        self.model = copy.deepcopy(model)

    def per_update_model(self, model_param):
        self.model.set_model_params(model_param)

    def get_train_num(self):
        return self.train_num

    def train(self, lr, epoch):
        self.model.train(self.train_data, lr = lr, epoch = epoch)
        w = self.model.get_model_params()
        return w

    def per_train(self, lr, epoch, per_layer):
        self.model.train(self.train_data, lr = lr, epoch = epoch)
        w = self.model.get_model_params(per_layer = per_layer)
        return w

    def my_personalization_train(self, lr, epoch, state_dict):
        self.model.my_personalization_train(self.train_data, lr = lr, epoch = epoch, state_dict = state_dict)
        w = self.model.get_model_params()
        return w

    def test(self, mode):
        if mode == 'train':
            dataset = self.train_data
        elif mode == 'test':
            dataset = self.test_data
        else:
            raise Exception('mode is train or test')

        metrics = self.model.test(dataset)
        return metrics

    def use_current_w_get_y(self, data):
        y = self.model.get_y(data = data)
        return y

    def get_label(self):
        each_label_num_dict = {}
        label_list = []
        for i in range(10):
            each_label_num_dict[i] = 0

        for batch_data in self.train_data:
            _, batch_y = batch_data
            for y in batch_y:
                y = int(y) #y是tensor类型, 所以要转化为int类型
                each_label_num_dict[y] += 1

        for each_label in each_label_num_dict:
            if each_label_num_dict[each_label] != 0:
                label_list.append(each_label)

        label_list = sorted(label_list)
        self.each_label_num_dict = each_label_num_dict
        self.label_list = label_list

    def get_hash_label(self):
        hash_label_list = []
        for label in self.label_list:
            hash_obj = SHA256.new()
            str_label = str(label)
            byte_label = bytes(str_label, 'utf-8')
            hash_obj.update(byte_label)
            hash_label_list.append(hash_obj.digest())
        self.hash_label_list = hash_label_list
        return hash_label_list

    def send(self, message = 'default'):
        if message == 'default':
            message = self.hash_label_list
        self.pipe_conn.send(message)

    def receive(self):
        return self.pipe_conn.recv()

    def get_intersection(self, received_list):
        j = 0
        flag = 0
        for i in range(len(received_list)):
            while j < len(self.hash_label_list):
                if received_list[i] == self.hash_label_list[j]:
                    self.intersection.append(self.label_list[j])
                    j += 1
                    flag = j
                    break
                else:
                    j += 1
            
            if flag > 0:
                j = flag
            else:
                j = 0
    
    def calculate_percentage(self):
        num = 0
        for i in self.intersection:
            num += self.each_label_num_dict[i]
        # p = num / self.total_num
        p = num / self.train_num  #在这里total_num就是train num
        return p

class C_Per_Client(Client):
    def __init__(self, client_idx, client_train_data, client_test_data, client_data_num, model, Perlocal_model) -> None:
        super(C_Per_Client, self).__init__(client_idx, client_train_data, client_test_data, client_data_num, model)
        self.Perlocal_model = Perlocal_model

    def Cper_train(self, state_dict, Cper_epoch, Cper_lr):
        self.Perlocal_model.Cper_train(train_data = self.train_data, Cper_epoch = Cper_epoch, Cper_lr = Cper_lr, state_dict = state_dict)

    def Cper_test(self, mode):
        if mode == 'train':
            dataset = self.train_data
        elif mode == 'test':
            dataset = self.test_data
        else:
            raise Exception('mode is train or test')

        metrics = self.Perlocal_model.test(dataset)
        return metrics